#!/usr/bin/env python

from threading import Thread

import numpy as np
import vtk
from scipy.spatial.transform import Rotation
from vtk.util.numpy_support import numpy_to_vtk, vtk_to_numpy

from .XmlParser import RobotTree
from .ZenLogger import ZenLogger

ZenLogger(console_log_level=ZenLogger.DEBUG)
logger = ZenLogger.get_logger(__name__)


def VtkTransformToNpArr(vtk_transform: vtk.vtkTransform):
    m: vtk.vtkMatrix4x4 = vtk_transform.GetMatrix()
    return np.array([m.GetElement(i, j) for i in range(4) for j in range(4)]).reshape(4, 4)


def VtkMatrix4x4ToNpArr(vtk_matrix4x4: vtk.vtkMatrix4x4):
    return np.array([vtk_matrix4x4.GetElement(i, j) for i in range(4) for j in range(4)]).reshape(4, 4)


def LoadObjMesh(mesh_file: str, renderer=None):
    reader = vtk.vtkOBJReader()
    reader.SetFileName(mesh_file)
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputConnection(reader.GetOutputPort())
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    if renderer:
        renderer.AddActor(actor)
    return actor


def LoadStlMesh(mesh_file: str, renderer=None):
    reader = vtk.vtkSTLReader()
    reader.SetFileName(mesh_file)
    mapper = vtk.vtkPolyDataMapper()
    mapper.SetInputConnection(reader.GetOutputPort())
    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    if renderer:
        renderer.AddActor(actor)
    return actor


def LoadDaeMesh(mesh_file: str, renderer=None):
    import trimesh
    dae_data = trimesh.exchange.dae.load_collada(mesh_file)
    geo = dae_data.get('geometry')
    asm_actor = vtk.vtkAssembly()
    rgba = [0.7, 0.7, 0.7, 1.0]
    if geo:
        for _, mesh_data in geo.items():
            trimesh.exchange.export.export_mesh(
                trimesh.Trimesh(**mesh_data), "./temp.stl")
            reader = vtk.vtkSTLReader()
            reader.SetFileName("./temp.stl")
            mapper = vtk.vtkPolyDataMapper()
            mapper.SetInputConnection(reader.GetOutputPort())
            mapper.Update()
            actor = vtk.vtkActor()
            actor.SetMapper(mapper)
            asm_actor.AddPart(actor)
    if renderer:
        renderer.AddActor(asm_actor)
    return asm_actor


def LoadMesh(mesh_file: str, renderer=None):
    actor = None
    if mesh_file.endswith(".stl"):
        actor = LoadStlMesh(mesh_file)
    elif mesh_file.endswith(".obj"):
        actor = LoadObjMesh(mesh_file)
    elif mesh_file.endswith(".dae"):
        actor = LoadDaeMesh(mesh_file)
    return actor


def LoadUrdf(render, urdf_path: str, show_edge=False, show_face=True):
    robot_tree = RobotTree(urdf_path)
    joints = robot_tree.GetJoints()
    links = robot_tree.GetLinks()
    pos = np.zeros(3)
    rot = np.zeros(3)
    # link_to_base = np.eye(4)
    actors = {}
    axes_actors = []
    transform1 = None
    for node in robot_tree.GetChain():
        link = links.get(node, False)
        if not link or link.get("visual_mesh", "") == "":
            continue
        f: str = link.get("visual_mesh", "")
        actor = LoadMesh(f)
        if not actor:
            return None, None
        try:
            rgba = link.get("color", (0.5, 0.5, 0.5, 1.0))
            prop = actor.GetProperty()
            prop.SetColor(*rgba[:3])
            prop.SetOpacity(rgba[3])
        except Exception:
            parts: vtk.vtkProp3DCollection = actor.GetParts()
            num = parts.GetNumberOfItems()
            for i in range(num):
                rgba = link.get("color", (0.5, 0.1, 0.8, 1.0))
                prop = parts.GetItemAsObject(i).GetProperty()
                prop.SetColor(*rgba[:3])
                prop.SetOpacity(rgba[3])
        render.AddActor(actor)
        actors[link['name']] = actor
        logger.info(f"link: {link['name']}")
        if link["parent"]:
            parent_joint_pos = joints[link["parent"]]["origin_xyz"]
            parent_joint_rpy = joints[link["parent"]]["origin_rpy"]
            # link_to_parent = np.eye(4)
            # link_to_parent[:3, 3] = parent_joint_pos
            # link_to_parent[:3, :3] = Rotation.from_euler(
            #     "zyx", parent_joint_rpy).as_matrix()
            # link_to_base = link_to_base @ link_to_parent
            transform2 = vtk.vtkTransform()
            transform2.Translate(parent_joint_pos)
            # TODO: check the follow code, rpy sequence
            transform2.RotateZ(parent_joint_rpy[2])
            transform2.RotateX(parent_joint_rpy[1])
            transform2.RotateY(parent_joint_rpy[0])
            # transform2.SetMatrix(link_to_parent.ravel())
            transform2.SetInput(transform1)
            actor.SetUserTransform(transform2)
            transform1 = transform2

            axes = vtk.vtkAxesActor()
            axes.AxisLabelsOff()
            axes.DragableOff()
            axes.SetShaftTypeToCylinder()
            axes.SetTotalLength((0.2, 0.2, 0.2))
            axes.SetUserTransform(transform2)
            render.AddActor(axes)
            axes_actors.append(axes)
    return actors, robot_tree


def TestVtkInteractor(urdf_path: str):
    from .Controller import Controller
    renderer = vtk.vtkRenderer()
    renderer.SetBackground(1.0, 1., 1.0)

    # Set initial camera position
    cam: vtk.vtkCamera = renderer.GetActiveCamera()
    cam.SetPosition(0, -5, 5)

    # Set global coordinate axes
    axes0 = vtk.vtkAxesActor()
    axes0.SetPosition(0, 0, 0)
    axes0.AxisLabelsOff()
    axes0.DragableOff()
    axes0.SetShaftTypeToCylinder()
    axes0.SetTotalLength((0.3, 0.3, 0.3))
    renderer.AddActor(axes0)

    # Load robot from urdf file
    actors, robot_tree = LoadUrdf(renderer, urdf_path)
    controller = Controller(actors, robot_tree)
    controller.StartServer()

    renwin = vtk.vtkRenderWindow()
    renwin.AddRenderer(renderer)
    renwin.SetWindowName("RobotZen")
    renwin.SetSize(800, 800)

    iren = vtk.vtkRenderWindowInteractor()
    iren.SetRenderWindow(renwin)
    style = vtk.vtkInteractorStyleTrackballCamera()
    iren.SetInteractorStyle(style)

    # Add left-bottom axes
    axes = vtk.vtkAxesActor()
    widget = vtk.vtkOrientationMarkerWidget()
    widget.SetOutlineColor(0.9300, 0.5700, 0.1300)
    widget.SetOrientationMarker(axes)
    widget.SetInteractor(iren)
    widget.SetEnabled(1)
    widget.InteractiveOff()

    iren.Initialize()
    iren.AddObserver(vtk.vtkCommand.KeyPressEvent, lambda *args,
                     **kwargs: print(f"args: {args}, kwargs: {kwargs}"))
    iren.AddObserver(vtk.vtkCommand.TimerEvent,
                     lambda ob, ev: controller._Update() and iren.Render())
    iren.AddObserver("LeftButtonPressEvent", lambda *args, **
                     kwargs: print(f"args: {args}, kwargs: {kwargs}"))
    iren.CreateRepeatingTimer(50)

    # ZenLogger(console_log_level=ZenLogger.INFO)
    # link1 = actors['link_1_s']
    # link2 = actors['link_2_l']
    # link3 = actors['link_3_u']
    # link4 = actors['link_4_r']
    # link5 = actors['link_5_b']
    # link6 = actors['link_6_t']
    # ut1 = link1.GetUserTransform()
    # ut2 = link2.GetUserTransform()
    # ut3 = link3.GetUserTransform()
    # ut4 = link4.GetUserTransform()
    # ut5 = link5.GetUserTransform()
    # ut6 = link6.GetUserTransform()

    iren.Start()


if __name__ == "__main__":
    urdf_path = "/home/xiaodaxia/Rvbust/Sources/RobotZenVispy/models/GP12/GP12.urdf"
    link_path = "/home/xiaodaxia/Rvbust/Sources/RobotZenVispy/models/GP12/visual/gp12_link_2_l.stl"
    TestVtkInteractor(urdf_path)
